import numpy as np
from data_process import *
from config_colab import Config
from tqdm import tqdm


label_dic = json.load(open(Config.label_dic_path))

reverse_label_dic = dict([(item[1], item[0])
                          for item in label_dic.items()])


def write_ensemble_submission(
        test_probs,
        text_file=Config.test_path,
        submission_path=None,
        is_combain=False):
    """Write an ensemble submission .

    Args:
        test_probs_files ([type]): [description]
        text_file ([type]): [description]
        submission_path ([type]): [description]
        is_combain (bool): [description]
    """

    if isinstance(test_probs, list):
        sum_probs = None

        for test_prob_f in tqdm(test_probs_files):

            test_prob_s = np.load(test_prob_f)

            if sum_probs is None:
                sum_probs = sum_probs+test_prob_s
    else:
        sum_probs = test_probs

    submission_list = build_submission(
        open(text_file).readlines(), sum_probs, is_combain)

    open(submission_path, mode="w").writelines(submission_list)


def build_submission(test_text, test_probs, is_combain):
    """Function that runs the submission function on test_x .

    Args:
        test_x ([type]): [description]
        test_probs ([type]): [description]

    Returns:
        [type]: [description]
    """

    probs_max = np.argmax(test_probs, axis=-1)

    submission_list = []

    for i, item in tqdm(enumerate(probs_max)):

        text = test_text[i].split("\x01")[1].strip()

        labels = [reverse_label_dic[idx] for idx in item[1:len(text)+1]]

        submission_list.append(
            test_text[i].strip() + "\u0001"+" ".join(labels)+"\n")
    if not is_combain:

        return submission_list

    submission_list_new = []

    for line in submission_list:

        line = post_process(line)
        submission_list_new.append(line)

    return submission_list_new
